package com.thqmcp.server.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.thqmcp.server.model.McpMessage;
import com.thqmcp.server.service.McpService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

/**
 * MCP WebSocket处理器
 * 处理基于WebSocket的MCP协议通信
 */
@Component
@RequiredArgsConstructor
@Slf4j
public class McpWebSocketHandler implements WebSocketHandler {

    private final McpService mcpService;
    private final ObjectMapper objectMapper = new ObjectMapper();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("WebSocket连接已建立: {}", session.getId());
        
        // 发送欢迎消息
        McpMessage welcomeMessage = new McpMessage();
        welcomeMessage.setJsonrpc("2.0");
        welcomeMessage.setType("notification");
        welcomeMessage.setMethod("server/ready");
        
        String message = objectMapper.writeValueAsString(welcomeMessage);
        session.sendMessage(new TextMessage(message));
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        log.info("收到WebSocket消息: {}", message.getPayload());
        
        try {
            if (message instanceof TextMessage textMessage) {
                String payload = textMessage.getPayload();
                McpMessage mcpMessage = objectMapper.readValue(payload, McpMessage.class);
                
                // 处理MCP消息
                McpMessage response = mcpService.handleMessage(mcpMessage);
                
                // 发送响应
                String responseJson = objectMapper.writeValueAsString(response);
                session.sendMessage(new TextMessage(responseJson));
            }
        } catch (Exception e) {
            log.error("处理WebSocket消息失败", e);
            
            // 发送错误响应
            McpMessage errorResponse = new McpMessage();
            errorResponse.setJsonrpc("2.0");
            errorResponse.setType("error");
            errorResponse.setError(new com.thqmcp.server.model.McpError(-32603, "内部错误", e.getMessage()));
            
            String errorJson = objectMapper.writeValueAsString(errorResponse);
            session.sendMessage(new TextMessage(errorJson));
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("WebSocket传输错误: {}", session.getId(), exception);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        log.info("WebSocket连接已关闭: {}, 状态: {}", session.getId(), closeStatus);
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
} 